import argparse
import multiprocessing
import os
import pickle
import sys
import warnings
import xml.etree.ElementTree as ET
from multiprocessing import Manager, Pipe
from threading import Thread

import cv2
import ipdb
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import scipy.io as sio
import spectral.io.envi as envi
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
from linformer import Linformer
from numpy import flip
from spectral import imshow, view_cube
from sshkeyboard import listen_keyboard
from torch.nn.parallel import DistributedDataParallel as DDP

from AutoGPU import autoGPU
from GAN_training_utils import DataResult, TrainProcess, setup_seed
# from AutoGPU import autoGPU
from models import (_1DCNN, _2DCNN, _3DCNN, _3DCNN_1DCNN, _3DCNN_AM, PURE2DCNN,
                    PURE3DCNN, PURE3DCNN_2AM, SAE, SAE_AM, DBDA_network,
                    HamidaEtAl, LeeEtAl, SSRN_network, _2dCNN, myknn, mysvm)
from myTrans2 import Generator
# from ViT import ViT
from NViT import ViT as NViT
from utils import DataPreProcess, myplot, plot, setpath, splitdata

# class ListenThread(Thread):
#     def __init__(self):
#         '''
#         :param func: 可调用的对象
#         :param args: 可调用对象的参数
#         '''
#         super().__init__()   # 不要忘记调用Thread的初始化方法
        

#     def run(self):
#         while True:
#             if parentPipe.recv() == 'plot':
#                 T_aux.model.load_state_dict(T.bestmodel)
#                 test_acc = T_aux.evaluate(T.test_loader, T_aux.test_result)
#                 print('{} set Accuracy:{:.2%}'.format('Test', test_acc))
#                 torch.save(T.bestmodel, resultpath + 'bestmodel.pth')
#                 myplot(processeddata, IMAGE, imagepath, T_aux.test_result, '_%s'%(test_acc))


class MyThread(Thread):
    def __init__(self):
        '''
        :param func: 可调用的对象
        :param args: 可调用对象的参数
        '''
        super().__init__()   # 不要忘记调用Thread的初始化方法
        

    def run(self):
        while True:
            if parentPipe.recv() == 'plot':
                dis_model_aux.model.load_state_dict(T.bestmodel)
                test_acc = dis_model_aux.evaluate(T.test_loader, result_aux)
                print('{} set Accuracy:{:.2%}'.format('Test', test_acc))
                print('saving model' + '.'*20)
                torch.save(T.bestmodel, resultpath + 'bestmodel.pth')
                myplot(processeddata, IMAGE, imagepath, result_aux, '_%s'%(test_acc))
           

if __name__ == '__main__':
    result_aux = DataResult()
    print('='*30)
    setup_seed(1993)
    warnings.filterwarnings("ignore")
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset', type=str,  default='../pathology/data/032370b-20x-roi2', metavar='D',
                        help='the dataset path you load')
    parser.add_argument('--trial_number', type=int,  default=1, metavar='T',
                        help='the time you do this trial')
    parser.add_argument('--train_number', type=int,  default=2000, metavar='NTr', 
                        help='number of training set')
    parser.add_argument('--valid_number', type=int,  default=1000, metavar='NVa',
                        help='number of valid set')
    parser.add_argument('--test_number', type=int,  default=1000, metavar='NTe',
                        help='number of test set')
    parser.add_argument('--patchsize', type=int,  default=9, metavar='P',
                        help='patchsize of data') 
    parser.add_argument('--modelname', type=str,  default='TransGan', metavar='P', help='which model to choose') 
    parser.add_argument('--gpu_ids', type=int,  default=-1, metavar='G',
                        help='which gpu to use')
    
    parser.add_argument('--local_rank', type=int, default=1) 
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 
    args = parser.parse_args()
    
    dataset = args.dataset
    NTr = args.train_number
    trialnumber = args.trial_number
    NTe = args.test_number
    NVa = args.valid_number
    patchsize = args.patchsize
    modelname = args.modelname


    rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])

    dark_ref = envi.open('%s.hdr'%dataset, '%s.raw'%dataset)
    dark_nparr = np.array(dark_ref.load())
    dark_nparr = flip(dark_nparr, 0)
    tree = ET.parse('%s.xml'%dataset)
    plt.figure()
    plt.imshow(dark_nparr[:,:,10])
    plt.savefig('test.jpg') 
        # get root element
    polygon = tree.getroot().findall('object/polygon')

    mask = np.zeros((dark_nparr.shape[0], dark_nparr.shape[1]), dtype="uint8")
    for p in polygon:
        x = p.findall('pt/x')
        y = p.findall('pt/y')
        x_coor = list(map(lambda x:int(x.text), x))
        y_coor = list(map(lambda y:int(y.text), y))
        c = []
        for x, y in zip(x_coor, y_coor):
            c.append([x, y])
        cor_xy = np.array(c)
        # cor_xy = np.hstack(mas(x_coor, y_coor))
        cv2.polylines(mask, np.int32([cor_xy]), 1, 1)
        cv2.fillPoly(mask, np.int32([cor_xy]), 1)


    resultpath, imagepath, datapath = setpath(dataset, trialnumber , NTr, NVa, NTe, modelname)
    IMAGE = np.flipud(dark_nparr)
    GND = mask
    spliteddata = splitdata(IMAGE, GND, datapath , trainnum=NTr, validnum=NVa, testnum=NTe)
    processeddata = DataPreProcess(IMAGE, patchsize, datapath, 1).processeddata
    parentPipe, childPipe = Pipe(True)
    dis_model = NViT().to('cuda')

    dis_model_aux = NViT().to('cuda')

    torch.distributed.init_process_group(backend='nccl')
    local_rank = torch.distributed.get_rank()
    # dis_model = DDP(dis_model, device_ids=[local_rank], output_device=local_rank)
    dis_model = DDP(dis_model, device_ids=[local_rank], output_device=local_rank)
    gen_model = Generator().to('cuda')
    # gen_model = DDP(gen_model, device_ids=[local_rank], output_device=local_rank)
    gen_model = DDP(gen_model, device_ids=[local_rank], output_device=local_rank)
    
    
    T = TrainProcess({'Gen':gen_model, 'Dis':dis_model},
                    processeddata,
                    './config.yaml',
                     childPipe,
                     writerpath=resultpath)    

    t = MyThread()
    t.setDaemon(True) 
    t.start()
    T.training_start()


   
    with open(resultpath + 'result.pkl', 'wb') as f:
        pickle.dump(T.test_result, f, pickle.HIGHEST_PROTOCOL)
    # myplot(processeddata, IMAGE, imagepath, T.test_result)
